import requests
import json

from redis import Redis

# redis_client = Redis(host='127.0.0.1', port=6379, db=0, decode_responses=True, password='vsoft@123&')
redis_client = Redis(host='127.0.0.1', port=6379, db=0, decode_responses=True)


# redis_client = Redis(host='10.206.60.14', port=6379, db=0, decode_responses=True)


def ollama_chat(prompt, assistant_message_id, model="gemma:7b"):
    redis_client.set(assistant_message_id, "start")
    # 设置键的失效时间为60秒
    redis_client.expire(assistant_message_id, 600)
    url = 'http://127.0.0.1:11434/api/generate'
    data = {
        # "model": "gemma:7b",
        "model": model,
        "prompt": prompt
    }
    headers = {'Content-Type': 'application/json'}
    response = requests.post(url, data=json.dumps(data), headers=headers, stream=True)
    full_response = ""
    for line in response.iter_lines():
        stop_stream = redis_client.get(assistant_message_id)
        if stop_stream == "end":
            break
        if line:
            res = json.loads(line.decode('utf-8'))
            print(res.get('response', None))
            full_response += res.get('response', None)
            yield json.dumps({'text': res.get('response', None)}) + "&&##&&"


def ollama_chat_all(prompt, assistant_message_id, model="gemma:7b"):
    redis_client.set(assistant_message_id, "start")
    # 设置键的失效时间为60秒
    redis_client.expire(assistant_message_id, 600)
    url = 'http://127.0.0.1:11434/api/generate'
    data = {
        # "model": "gemma:7b",
        "model": model,
        "prompt": prompt
    }
    headers = {'Content-Type': 'application/json'}
    response = requests.post(url, data=json.dumps(data), headers=headers, stream=True)
    full_response = ""
    for line in response.iter_lines():
        stop_stream = redis_client.get(assistant_message_id)
        if stop_stream == "end":
            break
        if line:
            res = json.loads(line.decode('utf-8'))
            print(res.get('response', None))
            full_response += res.get('response', None)
    return full_response


def fuzi_chat(model_suffix, user_text, assistant_message_id, history=None):
    url = 'http://127.0.0.1:5002/fuzi/chat'
    data = {
        "model_suffix": model_suffix,
        "user_text": user_text,
        "assistant_message_id": assistant_message_id,
        "history": history
    }

    headers = {'Content-Type': 'application/json'}
    response = requests.post(url, data=json.dumps(data), headers=headers, stream=True)
    for line in response.iter_lines():
        if line:
            # res = json.loads(line.decode('utf-8'))
            print(line)
            yield line


def fuzi_stop(assistant_message_id):
    url = 'http://10.206.60.14:5001/fuzi/stop'
    data = {
        "assistant_message_id": assistant_message_id
    }

    headers = {'Content-Type': 'application/json'}
    response = requests.post(url, data=json.dumps(data), headers=headers, stream=True)
